'''
KNN 插值
'''

import sys
sys.path.append(".")
sys.path.append('..')

from fancyimpute import MatrixFactorization
import numpy as np


from impute_method.impute_base import ImputeBase


class MFImpute(ImputeBase):
    def __init__(self, **kwargs):
        super(MFImpute, self).__init__(**kwargs)
        # self.kwargs = kwargs
        self.method = 'matrixFactorization'

    def impute(self):
        masked_data = self.masked_data[:, 0, :]  # 转化为二维 [n_samples, feature_dim]
        masked_data[masked_data == 0] = np.nan
        imputed_data = MatrixFactorization(max_iters=self.kwargs['max_iters'], rank=self.kwargs['rank'],
                                           shrinkage_value=self.kwargs['shrinkage_value'])\
            .fit_transform(masked_data)
        return self.devide_window(imputed_data)


def main(seq_len, max_iters, rank, shrinkage_value):
    kwargs = dict()
    kwargs['data_name'] = 'water'
    kwargs['data_type'] = 'masked'
    kwargs['indicators'] = 'WATER_TEMPERATURE,PH_VALUE,TOTAL_NITROGEN,DISSOLVED_OXYGEN'
    kwargs['masked_indicator'] = ''
    kwargs['seq_len'] = seq_len
    # 插值方法的参数
    kwargs['max_iters'] = max_iters
    kwargs['rank'] = rank
    kwargs['shrinkage_value'] = shrinkage_value

    mf_impute = MFImpute(**kwargs)
    mf_impute.metirc()


if __name__ == '__main__':
    main(24, 200, 30, 0.3)


